package com.sh.sh_ai_agent.agent;

import cn.hutool.core.util.StrUtil;
import com.itextpdf.styledxmlparser.jsoup.internal.StringUtil;
import com.sh.sh_ai_agent.advisor.MyLoggerAdvisor;
import com.sh.sh_ai_agent.agent.model.AgentState;
import jakarta.annotation.Resource;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * 抽象基础代理类，用于管理代理状态和执行流程。
 * 提供状态转换、内存管理和基于步骤的执行循环的基础功能。
 * 子类必须实现step方法。
 */
@Data
@Slf4j
public abstract class BaseAgent {

    // 核心属性
    private String name;
    // 提示词
    private String systemPrompt;
    private String nextStepPrompt;
    // 代理状态
    private AgentState state = AgentState.IDLE;

    // 步骤
    private int currentStep = 0;
    private int maxSteps = 10;
    // 大模型
    private ChatClient chatClient;

    // 上下文记忆
    private List<Message> messageList = new ArrayList<>();
    // 执行流程

    /**
     * 运行代理，执行流程。
     */
    public String run(String userPrompt) {
        if (this.state != AgentState.IDLE) {
            throw new RuntimeException("cannot run agent in state " + this.state);
        }
        if (StrUtil.isBlank(userPrompt)) {
            throw new RuntimeException("userPrompt is blank");
        }
        state = AgentState.RUNNING;
        messageList.add(new UserMessage(userPrompt));
        // 结果列表
        List<String> results = new ArrayList<>();

        try {
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Running step " + stepNumber + "/" + maxSteps);
                String stepResult = step();
                String result = "Step " + currentStep + ": " + stepResult;
                results.add(result);
            }
            if (currentStep >= maxSteps) {
                state = AgentState.FINISHED;
                log.warn("reached max steps");
                results.add("reached max steps");
            }
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("error running agent", e);
            return "执行错误" + e.getMessage();
        } finally {
            cleanUp();
        }
    }

    /**
     * 执行单个步骤。
     */
    public abstract String step();


    /**
     * 清理资源
     */
    public void cleanUp() {
        // 清理资源，例如关闭连接等。
    }



    /**
     * 运行代理（流式输出）
     *
     * @param userPrompt 用户提示词
     * @return SseEmitter实例
     */
    public SseEmitter runStream(String userPrompt) {
        // 创建SseEmitter，设置较长的超时时间
        SseEmitter emitter = new SseEmitter(300000L); // 5分钟超时

        // 使用线程异步处理，避免阻塞主线程
        CompletableFuture.runAsync(() -> {
            try {
                if (this.state != AgentState.IDLE) {
                    emitter.send("错误：无法从状态运行代理: " + this.state);
                    emitter.complete();
                    return;
                }
                if (StringUtil.isBlank(userPrompt)) {
                    emitter.send("错误：不能使用空提示词运行代理");
                    emitter.complete();
                    return;
                }

                // 更改状态
                state = AgentState.RUNNING;
                // 记录消息上下文
                messageList.add(new UserMessage(userPrompt));

                try {
                    for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                        int stepNumber = i + 1;
                        currentStep = stepNumber;
                        log.info("Executing step " + stepNumber + "/" + maxSteps);

                        // 单步执行
                        String stepResult = step();
                        if (stepResult.equals("没有工具需要调用")) {
                            // 使用大模型进行自然对话（记忆式），返回 Flux<String>
                            Flux<String> stringFlux = chatClient
                                    .prompt()
                                    .user(userPrompt)
                                    .advisors(spec -> spec
                                            .param(CHAT_MEMORY_CONVERSATION_ID_KEY, 1)
                                            .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10))
                                    .advisors(new MyLoggerAdvisor())
                                    .stream()
                                    .content();// Flux<String> 流式输出
                            // 将流式内容逐段发送给 SSE emitter
                            stringFlux
                                    .doOnNext(token -> {
                                        try {
                                            emitter.send("回复: " + token);
                                        } catch (IOException e) {
                                            log.error("发送流式内容失败", e);
                                        }
                                    })
                                    .doOnComplete(() -> {
                                        try {
                                            emitter.send("对话结束。");
                                            emitter.complete();
                                        } catch (IOException e) {
                                            log.error("发送完成失败", e);
                                        }
                                    })
                                    .doOnError(err -> {
                                        try {
                                            emitter.send("对话出错: " + err.getMessage());
                                            emitter.completeWithError(err);
                                        } catch (IOException e) {
                                            emitter.completeWithError(e);
                                        }
                                    })
                                    .subscribe(); // 👈 关键：必须订阅才会触发执行
                            return ; // 避免进入 stepResult 的 emitter.send(result)
                        }

                        String result = "Step " + stepNumber + ": " + stepResult;

                        // 发送每一步的结果
                        emitter.send(result);
                    }
                    // 检查是否超出步骤限制
                    if (currentStep >= maxSteps) {
                        state = AgentState.FINISHED;
                        emitter.send("执行结束: 达到最大步骤 (" + maxSteps + ")");
                    }
                    // 正常完成
                    emitter.complete();
                } catch (Exception e) {
                    state = AgentState.ERROR;
                    log.error("执行智能体失败", e);
                    try {
                        emitter.send("执行错误: " + e.getMessage());
                        emitter.complete();
                    } catch (Exception ex) {
                        emitter.completeWithError(ex);
                    }
                } finally {
                    // 清理资源
                    this.cleanUp();
                }
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        // 设置超时和完成回调
        emitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanUp();
            log.warn("SSE connection timed out");
        });

        emitter.onCompletion(() -> {
            if (this.state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            this.cleanUp();
            log.info("SSE connection completed");
        });

        return emitter;
    }




}
